from typing import Callable

from ril import DQNRulePolicy, PGRulePolicy, SACRulePolicy, DiscreteSACRulePolicy, A2CRulePolicy, PPORulePolicy


def common_policy_factory(policy: str, rule_fn: Callable = None):
    def _policy_factory(*args, **kargs):
        if policy == 'dqn':
            return DQNRulePolicy(rule_fn, *args, **kargs, is_double=False)
        if policy in ['double dqn', 'ddqn', 'dueling dqn']:
            return DQNRulePolicy(rule_fn, *args, **kargs, is_double=True)
        if policy == 'pg':
            return PGRulePolicy(rule_fn, *args, **kargs)
        if policy == 'sac':
            return SACRulePolicy(rule_fn, *args, **kargs)
        if policy == 'discrete sac':
            return DiscreteSACRulePolicy(rule_fn, *args, **kargs)
        if policy == 'a2c':
            return A2CRulePolicy(rule_fn, *args, **kargs)
        if policy == 'ppo':
            return PPORulePolicy(rule_fn, *args, **kargs)
        else:
            raise Exception('unexpected agent: %s' % policy)

    return _policy_factory
